import numpy as np
import matplotlib.pyplot as plt

def plot_comparison_subplot(ax, total_file, sa_file, tdma_file, x_lable):
    # 加载数据
    total_data = np.loadtxt(total_file)
    sa_data = np.loadtxt(sa_file)
    tdma_data = np.loadtxt(tdma_file)

    # 绘制吞吐量曲线

    sa_line, = ax.plot(sa_data, color='orange', lw=1, label='Only SA')
    tdma_line, = ax.plot(tdma_data, color='green', lw=1, label='Only TDMA')
    total_line, = ax.plot(total_data, color='#236B8E', lw=1.2, label='Total Throughput')

    # 设置图例、坐标轴和标题
    ax.legend(handles=[total_line, sa_line, tdma_line])
    ax.set_xlim((0, 10005))
    ax.set_ylim((0, 1))
    #ax.set_xlabel('Time Slot')
    ax.set_xlabel(x_lable)
    ax.set_ylabel('Throughput')
    #ax.set_title(title)

def main():
    # 创建画布和子图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

    # 绘制第一个子图 4.2a
    plot_comparison_subplot(ax1,
                            'rewards/total_111.txt',
                            'rewards/only_SA_111.txt',
                            'rewards/only_TDMA_111.txt',
                            'Time Slot\na delay∈(1,5) |U|=3')

    # 绘制第二个子图 4.2b
    plot_comparison_subplot(ax2,
                            'rewards/total_122.txt',
                            'rewards/only_SA_122.txt',
                            'rewards/only_TDMA_122.txt',
                            'Time Slot\nb delay∈(1,5) |U|=5')

    # 调整布局
    plt.tight_layout()

    # 显示图形
    plt.show()

if __name__ == "__main__":
    main()
